﻿// ref RFC 6455
using System;
using System.IO;
using System.Net.Security;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using System.Text;
using System.Threading;
using yunio.helpers;

namespace yunio._req_ws
{
    internal class ws_client : IYunioWsClient
    {
        private TcpClient _tcpClient;
        private NetworkStream _rawStream; // 原始 TCP 流（未加密）
        private SslStream _sslStream;     // TLS 加密流（仅 wss 使用）
        private Stream _activeStream;     // 当前使用的流（_rawStream 或 _sslStream）
        private bool _isConnected;
        private string _secWebSocketKey;
        private string _host;                   // 服务端主机（如 localhost:1111）
        private string _path;                   // 完整路径（含查询参数，如 /funca?t=1）
        private bool _isWss;                    // 是否为 wss 协议

        public event YunioCallbackFunction OnConnected;
        public event YunioCallbackFunction OnDisconnected;
        public event YunioOnMessage OnMessage;

        public string uri { get; internal set; }

        public bool IsConnected => _isConnected;
        private byte[] buffer;
        internal MemoryStream ms_cache;
        internal IAsyncResult read_handle;

        /// <summary>
        /// 初始化客户端（自动解析 URL）
        /// </summary>
        /// <param name="url">WebSocket URL（如 ws://localhost:1111/funca?t=1 或 wss://...）</param>
        /// <param name="heartbeatInterval">心跳间隔（毫秒）</param>
        /// <param name="heartbeatTimeout">心跳超时时间（毫秒）</param>
        public ws_client(string uri, int cache_size = 4096)
        {
            buffer =new byte[cache_size];
            ms_cache = new MemoryStream();
            this.uri = uri;
            ParseUrl(uri); // 解析 URL 中的协议、主机、路径等
        }

        /// <summary>
        /// 连接到 WebSocket 服务端
        /// </summary>
        public bool Connect()
        {
            try
            {
                // 建立 TCP 连接
                _tcpClient = new TcpClient();
                _tcpClient.Connect(_host.Split(':')[0], int.Parse(_host.Split(':')[1]));
                _rawStream = _tcpClient.GetStream();
                _isConnected = true;

                // 若为 wss，启用 TLS 加密
                if (_isWss)
                {
                    _sslStream = new SslStream(_rawStream, false, ValidateServerCertificate);
                    _sslStream.AuthenticateAsClient(_host); // 同步 TLS 握手
                    _activeStream = _sslStream;
                }
                else
                {
                    _activeStream = _rawStream;
                }

                // 生成握手 Key
                var rng = new RNGCryptoServiceProvider();
                byte[] keyBytes = new byte[16];
                rng.GetBytes(keyBytes);
                _secWebSocketKey = Convert.ToBase64String(keyBytes);

                // 发送 WebSocket 握手请求
                string handshake = BuildHandshakeRequest();
                byte[] handshakeBytes = Encoding.UTF8.GetBytes(handshake);
                _activeStream.Write(handshakeBytes, 0, handshakeBytes.Length);

                // 验证握手响应
                if (!ValidateHandshake())
                    throw new Exception("WebSocket 握手失败");

                OnConnected?.Invoke(); // 触发连接成功事件
                return true;
            }
            catch (Exception ex)
            {
                Console.WriteLine($"连接错误：{ex.Message}");
                Disconnect();
            }
            return false;
        }

        /// <summary>
        /// 发送文本消息
        /// </summary>
        public void Send(string message)
        {
            if (!_isConnected) throw new InvalidOperationException("未连接");
            byte[] frame = BuildTextFrame(message);
            _activeStream.Write(frame, 0, frame.Length); // 统一通过当前流发送
        }

        /// <summary>
        /// 断开连接
        /// </summary>
        public void Disconnect()
        {
            if (_tcpClient?.Connected ?? false)
            {
                _isConnected = false;

                // 关闭 TLS 流（若有）
                if (_sslStream != null)
                {
                    _sslStream.Close();
                    _sslStream = null;
                }

                // 关闭原始流和 TCP 连接
                _rawStream?.Close();
                _tcpClient?.Close();

                OnDisconnected?.Invoke();
            }
        }
        private byte[] BuildTextFrame(string message)
        {
            // 1. 将文本消息转换为 UTF-8 字节数组（原始负载）
            byte[] payload = Encoding.UTF8.GetBytes(message);
            int payloadLength = payload.Length;

            // 2. 创建内存流，用于组装帧数据
            using (MemoryStream ms = new MemoryStream())
            {
                // 3. 写入帧头前 2 字节（FIN=1，Opcode=0x01）
                //    FIN=1（二进制 1000 0000） | Opcode=0x01（二进制 0000 0001） → 0x81
                ms.WriteByte(0x81);

                // 4. 处理负载长度（简化版：仅支持 ≤125 字节的消息）
                if (payloadLength <= 125)
                {
                    // 直接写入 7 bits 长度（负载长度 ≤125）
                    ms.WriteByte((byte)payloadLength);
                }
                else
                {
                    // （示例未处理长消息，实际需扩展）
                    // 若长度 126：写入 126（0x7E），后跟 2 字节的负载长度（大端序）
                    // 若长度 127：写入 127（0x7F），后跟 8 字节的负载长度（大端序）
                    throw new NotSupportedException("仅支持短消息（负载长度 ≤125 字节）");
                }

                // 5. 生成 4 字节随机掩码（客户端必须使用掩码）
                byte[] mask = new byte[4];
                new RNGCryptoServiceProvider().GetBytes(mask); // 加密随机数生成器
                ms.Write(mask, 0, 4); // 写入掩码到帧中

                // 6. 应用掩码到原始负载（异或操作）
                for (int i = 0; i < payloadLength; i++)
                {
                    payload[i] ^= mask[i % 4]; // 每个字节与掩码的第 i%4 位异或
                }

                // 7. 写入掩码后的负载数据
                ms.Write(payload, 0, payloadLength);

                // 8. 返回完整的帧字节数组
                return ms.ToArray();
            }
        }

        /// <summary>
        /// 解析 WebSocket URL（提取协议、主机、路径等）
        /// </summary>
        private void ParseUrl(string url)
        {
            if (!url.StartsWith("ws://", StringComparison.OrdinalIgnoreCase) &&
                !url.StartsWith("wss://", StringComparison.OrdinalIgnoreCase))
                throw new ArgumentException("无效的 WebSocket URL（需以 ws:// 或 wss:// 开头）");

            _isWss = url.StartsWith("wss://", StringComparison.OrdinalIgnoreCase);

            // 提取主机和端口（格式：host:port 或 host）
            int protocolEndIndex = _isWss ? 6 : 5; // "wss://" 或 "ws://" 的长度
            string hostPortPart = url.Substring(protocolEndIndex);
            int pathStartIndex = hostPortPart.IndexOf('/');
            if (pathStartIndex == -1)
            {
                _host = hostPortPart;
                _path = "/";
            }
            else
            {
                _host = hostPortPart.Substring(0, pathStartIndex);
                _path = hostPortPart.Substring(pathStartIndex);

                // 补充默认端口（ws 默认 80，wss 默认 443）
                if (!_host.Contains(":"))
                {
                    _host += _isWss ? ":443" : ":80";
                }
            }
        }

        /// <summary>
        /// 构造 WebSocket 握手请求
        /// </summary>
        private string BuildHandshakeRequest()
        {
            return $@"GET {uri} HTTP/1.1
Host: {_host}
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: {_secWebSocketKey}
Sec-WebSocket-Version: 13

"; // 空行结束头
        }

        /// <summary>
        /// 验证服务端握手响应
        /// </summary>
        private bool ValidateHandshake()
        {
            byte[] buffer = new byte[4096];
            int read = _activeStream.Read(buffer, 0, buffer.Length);
            string response = Encoding.UTF8.GetString(buffer, 0, read);

            // 检查状态码和 Accept 头
            return response.Contains("101 Switching Protocols") &&
                   response.Contains($"Sec-WebSocket-Accept: {ComputeAccept(_secWebSocketKey)}");
        }

        /// <summary>
        /// 计算服务端应返回的 Sec-WebSocket-Accept
        /// </summary>
        private string ComputeAccept(string key)
        {
            using (var sha1 = SHA1.Create())
            {
                byte[] hash = sha1.ComputeHash(Encoding.UTF8.GetBytes(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"));
                return Convert.ToBase64String(hash);
            }
        }
        /// <summary>
        /// 验证服务端证书（可自定义逻辑，如忽略自签名证书）
        /// </summary>
        private bool ValidateServerCertificate(object sender, X509Certificate certificate, X509Chain chain, SslPolicyErrors sslPolicyErrors)
        {
            if (sslPolicyErrors != SslPolicyErrors.None)
            {
                return true; // 忽略所有证书错误（仅测试用！）
            }
            return true;
        }

        /// <summary>
        /// 启动接收循环（后台线程）
        /// </summary>
        public void StartReceiveLoop()
        {
            new Thread(ReceiveLoop) { IsBackground = true }.Start();
        }
        private void read_callback(IAsyncResult ar)
        {
            var size = _activeStream.EndRead(ar);
            ms_cache.Write(buffer, 0, size);
            ms_cache.Flush();
            if (ms_cache.Length == 0)
            {
                Disconnect();
                return; // 连接关闭
            }
            // 解析帧（支持文本和 Pong 帧）
            while (has_fullframe(ms_cache))
            {
                var msg = ParseDataFrame(ms_cache);

                if (msg != null)
                    try { OnMessage?.Invoke(msg); }
                    catch { }
            }
            ms_cache.Position = ms_cache.Length;
            read_handle = _activeStream.BeginRead(buffer, 0, buffer.Length, read_callback, null);
        }
        private void ReceiveLoop()
        {
            if (_isConnected && read_handle == null) read_handle = _activeStream.BeginRead(buffer, 0, buffer.Length, read_callback, null);
            try { while (_isConnected) if (read_handle != null) read_handle.AsyncWaitHandle.WaitOne(); }
            catch (Exception ex)
            {
                Console.WriteLine($"接收错误：{ex.Message}");
                Disconnect();
            }
        }
        /// <summary>
        /// 解析 WebSocket 数据帧（文本或 Pong）
        /// </summary>
        private YunioData ParseDataFrame(MemoryStream ms)
        {
            ms.Position = 0;
            int b1 = ms.ReadByte();
            if (b1 == -1) return null;

            int opcode = b1 & 0x0F;
            if (opcode == 0) opcode = 0x01;

            switch (opcode)
            {
                case 0x01: // 文本帧
                    if (!is_final_frame(ms))
                    {
                        ParseBinaryFrame(ms);
                        return null;
                    }
                    else
                    {
                        ParseBinaryFrame(ms);
                        if (binary.Length == 0) return null;
                        var text = Encoding.UTF8.GetString(binary.ToArray());

                        binary.Position = 0; // 清空二进制缓存
                        binary.SetLength(0);
                        return new YunioData(text);
                    }
                case 0x02: // 二进制帧
                    if (!is_final_frame(ms))
                    {
                        ParseBinaryFrame(ms);
                        return null;
                    }
                    else
                    {
                        ParseBinaryFrame(ms);

                        if (binary.Length == 0) return null;
                        var mem = new MemoryStream();
                        binary.Position = 0;
                        helper.ReadAllToMemory(binary, mem);
                        mem.Position = 0;
                        
                        binary.Position = 0; // 清空二进制缓存
                        binary.SetLength(0);
                        return new YunioData(mem);
                    }
                case 0x08: // 关闭通知
                    ms_cache.Position = 0;
                    ms_cache.SetLength(0);
                    return null;
                case 0x09: // ping 帧
                case 0x0A: // Pong 帧

                // 保留帧
                case 0x03:
                case 0x04:
                case 0x05:
                case 0x06:
                case 0x07:
                case 0x0B:
                case 0x0C:
                case 0x0D:
                case 0x0E:
                case 0x0F:
                    var length = ms_cache.Length - (read_datalen(ms_cache) + read_lenblock_size(ms_cache) + 2);
                    if (length < 0) length = 0;
                    ms_cache.Position = length;
                    ms_cache.SetLength(length);
                    break;
                default:
                    ms_cache.Position = ms_cache.Length - 2;
                    ms_cache.SetLength(ms_cache.Position);
                    return null; // 其他帧类型暂不处理
            }
            return null;
        }
        MemoryStream binary = new MemoryStream();
        /**
0 1 2 3 4 5 6 7    0 1 2 3 4 5 6 7  0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7
+-+-+-+-+-------+  +-+-------------+ +-----------------------------+
|F|R|R|R| OP    |  |M| LENGTH      |   Extended payload length
|I|S|S|S| CODE  |  |A|             |  （if LENGTH=126）
|N|V|V|V|       |  |S|             |
| |1|2|3|       |  |K|             |
+-+-+-+-+-------+  +-+-------------+
|                      Extended payload length（if LENGTH=127）
+                                  +-------------------------------
|      Extended payload length     | Masking-key，if Mask set to 1
+----------------------------------+-------------------------------
|   Masking-key                    |       Data
+----------------------------------+-------------------------------
|                                Data
+----------------------------------+-------------------------------
        */
        private bool has_fullframe(MemoryStream ms)
        {
            if (ms.Length < 2) return false;
            ms.Position = 0;
            switch (ms.ReadByte() & 0x0f)
            {
                case 0x08: //close
                case 0x09: //ping
                case 0x0A: //pong
                    return true;
                case 0x00: // 分片，延续帧
                case 0x01: // 文本
                case 0x02: // 二进制
                default:
                    var lensize = read_lenblock_size(ms);
                    if (ms.Length < lensize + 2) return false;
                    if (ms.Length < read_datalen(ms) + lensize + 2) return false;
                    return true;
            }
        }
        internal int read_lenblock_size(Stream buf)
        {
            buf.Position = 1;
            var b2 = buf.ReadByte();
            int payloadLen = b2 & 0x7F;

            switch (payloadLen)
            {
                case 127: return 8;
                case 126: return 2;
                default: return 0;
            }
        }
        internal int read_datalen(Stream buf)
        {
            switch (read_lenblock_size(buf))
            {
                case 8:
                    buf.Position = 2;
                    var len = (buf.ReadByte() << 56) | (buf.ReadByte() << 48) | (buf.ReadByte() << 40) | (buf.ReadByte() << 32) | (buf.ReadByte() << 24) | (buf.ReadByte() << 16) | (buf.ReadByte() << 8) | buf.ReadByte();
                    if (len < 0) return 0;
                    return len;
                case 2:
                    buf.Position = 2;
                    var len2 = (buf.ReadByte() << 8) | buf.ReadByte();
                    if (len2 < 0) return 0;
                    return len2;
                case 0:
                    buf.Position = 1;
                    return buf.ReadByte() & 0x7F;
            }
            return 0;
        }
        internal bool is_final_frame(Stream buf)
        {
            buf.Position = 0;
            return 0 != (buf.ReadByte() & 0x80) ;
        }
        internal void ParseBinaryFrame(MemoryStream ms)
        {
            var payloadLen = read_datalen(ms);
            var pos2 = read_lenblock_size(ms) + 2;
            if (payloadLen > 0) binary.Write(ms.GetBuffer(), pos2, payloadLen);
            pos2 += payloadLen;

            var length = ms.Length - pos2;
            if (length > 0)
            {
                ms_cache.Position = 0;
                ms_cache.Write(ms.GetBuffer(), pos2, (int)length);
                ms_cache.Flush();
                ms_cache.SetLength(length);
            }
            else
            {
                ms_cache.Position = 0;
                ms_cache.SetLength(0);
            }
        }
        public void Dispose()
        {
            this.buffer = null;
            this.binary.Dispose();
            this.binary = null;
            this.Disconnect();
        }
    }
}
